#include "core/tensor.h"

#include <map>
#include <cstddef>
#include <cassert>

namespace llm_core
{

size_t Tensor::GetElnBytes()
{
    static std::map<DataType, size_t> type_size = {
        {Tensor::DataType::DATA_TYPE_INT8, 1},
        {Tensor::DataType::DATA_TYPE_UINT8, 1},
        {Tensor::DataType::DATA_TYPE_INT16, 2},
        {Tensor::DataType::DATA_TYPE_UINT16, 2},
        {Tensor::DataType::DATA_TYPE_INT32, 4},
        {Tensor::DataType::DATA_TYPE_UINT32, 4},
        {Tensor::DataType::DATA_TYPE_FP16, 2},
        {Tensor::DataType::DATA_TYPE_FP32, 4},
    };

    assert(type_size.size() == Tensor::DataType::DATA_TYPE_COUNT);

    return type_size[m_params.m_data_type];
}

} // namespace llm_core